# 使用PyQtGraph的示例
import datetime

import tushare as ts
from PyQt5 import QtCore, QtWidgets, QtGui
from matplotlib.dates import date2num
from QtUI07 import mapping
import sys
import pandas as pd
import pyqtgraph as pg


class MainWindow(object):
    def __init__(self):
        app = QtWidgets.QApplication(sys.argv)
        MainWindow = QtWidgets.QMainWindow()
        self.ui = mapping.Ui_MainWindow()
        self.ui.setupUi(MainWindow)

        self.ui.verticalLayout.addWidget(chart())

        MainWindow.show()
        sys.exit(app.exec_())


# 根据PyQtGraph的示例，我们新建一个绘制K线图的类CandlestickItem
class CandlestickItem(pg.GraphicsObject):
    def __init__(self, data):
        pg.GraphicsObject.__init__(self)
        self.data = data  ## data must have fields: time, open, close, min, max
        self.generatePicture()

    def generatePicture(self):
        self.picture = QtGui.QPicture()
        p = QtGui.QPainter(self.picture)
        p.setPen(pg.mkPen('w'))
        w = (self.data[1][0] - self.data[0][0]) / 3.
        for (t, open, close, min, max) in self.data:
            p.drawLine(QtCore.QPointF(t, min), QtCore.QPointF(t, max))
            if open > close:
                p.setBrush(pg.mkBrush('g'))
            else:
                p.setBrush(pg.mkBrush('r'))
            p.drawRect(QtCore.QRectF(t - w, open, w * 2, close - open))
        p.end()

    def paint(self, p, *args):
        p.drawPicture(0, 0, self.picture)

    def boundingRect(self):
        return QtCore.QRectF(self.picture.boundingRect())


# 创建一个函数chart()，用于通过tushare模块获取股票的历史数据，并将数据传入之前定义好的用于生成K线图的类
def chart():
    hist_data = ts.get_hist_data('600519', start='2017-05-01', end='2017-11-24')
    data_list = []
    axis = []
    for dates, row in hist_data.iterrows():
        # 将时间转换为数字
        date_time = datetime.datetime.strptime(dates, '%Y-%m-%d')
        t = date2num(date_time)
        # t = dict(enumerate(datetime))
        open, high, close, low = row[:4]
        datas = (t, open, close, low, high)
        data_list.append(datas)
        axis.append(t)
    # print(axis)
    axis_dict = dict(enumerate(axis))
    item = CandlestickItem(data_list)
    plt = pg.PlotWidget()
    # print(plt.getAxis('bottom'))
    plt.addItem(item)
    plt.showGrid(x=True, y=True)
    return plt


if __name__ == "__main__":
    MainWindow()
